package firewall

import (
	"azfirewall-gateway/models"
	"azfirewall-gateway/utils"
	"crypto/sha256"
	"fmt"
	"net/http"
	"sync"
	"sync/atomic"
	"time"
)

type CC struct {
	// 防火墙CC主记录，设计为一个站实例化一个，减少锁竞争
	counts      sync.Map // map[token:string]*models.ClientStat
	policy      *models.CCPolicy
	ticker      *time.Ticker
	newTickerCh chan time.Duration
	isRunning   bool
}

func NewCC(policy *models.CCPolicy) *CC {
	utils.MainLog.Debug("【CC防护进程】CC攻击被启用，进程已注入")
	if policy.IntervalSeconds.Nanoseconds() == 0 {
		utils.MainLog.Debug("【CC防护进程】传入0秒ticker")
		return nil
	}
	c := CC{
		counts:      sync.Map{},
		policy:      policy,
		ticker:      time.NewTicker(policy.IntervalSeconds * time.Second),
		newTickerCh: make(chan time.Duration),
		isRunning:   true,
	}
	go c.checker()
	return &c
}

func (Self *CC) checker() {
	utils.MainLog.Debug("【CC Checker】\"检测\"方法被调用")
	// 单开goroutine根据周期检查
	for {
		if !Self.isRunning {
			// 根据此条件关闭goroutine
			return
		}
		select {
		case newDuration, ok := <-Self.newTickerCh:
			if ok {
				Self.ticker = time.NewTicker(newDuration)
				continue
				// 动态修改检查周期
			}
			return
		case <-Self.ticker.C:
			Self.check()
		}
	}
}

func (Self *CC) check() {
	utils.MainLog.Debug("【CC Checker】抄的\"检测\"方法被调用")
	// 抄的没动
	Self.counts.Range(func(key, value interface{}) bool {
		clientID := key.(string)
		stat := value.(*models.ClientStat)
		//fmt.Println("CCAttackTick:", appID, clientID, stat)
		if stat.IsBlackIP == true {
			stat.RemainSeconds -= Self.policy.IntervalSeconds
			if stat.RemainSeconds <= 0 {
				Self.counts.Delete(clientID)
			}
		} else if stat.Count >= Self.policy.MaxCount {
			stat.Count = 0
			stat.IsBlackIP = true
			stat.RemainSeconds = Self.policy.BlockSeconds
		} else {
			Self.counts.Delete(clientID)
		}
		return true
	})
}

func (Self *CC) IdentifyCCAttack(r *http.Request, domainName string, clientIP string) bool {
	utils.MainLog.Debug("【CC Checker】CC攻击检测计次调用")
	preHashContent := clientIP
	preHashContent += domainName
	ua := r.Header.Get("User-Agent")
	preHashContent += ua
	cookie := r.Header.Get("Cookie")
	preHashContent += cookie
	clientID := SHA256Hash(preHashContent)
	clientIDStat, _ := Self.counts.LoadOrStore(clientID, &models.ClientStat{Count: 0, IsBlackIP: false, RemainSeconds: 0})
	clientStat := clientIDStat.(*models.ClientStat)
	// clientStat 对象为指针，可以直接修改，但是可能会存在并发问题，所以使用原子操作计数
	if clientStat.IsBlackIP == true {
		if atomic.LoadInt64(&clientStat.Count) == 0 {
			atomic.AddInt64(&clientStat.Count, 1)
		}
		return true
	}
	atomic.AddInt64(&clientStat.Count, 1)
	return false
}

func (Self *CC) UpdatePolicy(newPolicy *models.CCPolicy) {
	Self.policy = newPolicy
	Self.newTickerCh <- newPolicy.IntervalSeconds
	//
}

func (Self *CC) Stop() {
	Self.isRunning = false
}

func SHA256Hash(plaintext string) string {
	hash := sha256.New()
	hash.Write([]byte(plaintext))
	result := fmt.Sprintf("%x", hash.Sum(nil))
	return result
}

type CCGroup struct {
	group sync.Map // map[host:string]*CC
	total uint32
}

func NewCCGroup() *CCGroup {
	return &CCGroup{}
}

func (Self *CCGroup) LoadOrStore(host string, policy *models.CCPolicy) (*CC, bool) {
	cc, ok := Self.group.Load(host)
	if !ok {
		// LoadOrStore 返回false意思为存储
		newCC := NewCC(policy)
		cc, ok = Self.group.LoadOrStore(host, newCC)
		// 这里进行原子性的计数，采用LoadOrStore，为false意思为插入成功，所以进行总数加1
		// 为true时已经有别的goroutine进行了插入，所以这里放弃刚才新建的对象，直接返回已经插入的
		if !ok {
			atomic.AddUint32(&Self.total, 1)
		} else {
			newCC.Stop()
		}
		return cc.(*CC), false
	}
	return cc.(*CC), true
}

func (Self *CCGroup) Delete(host string) {
	cc, ok := Self.group.Load(host)
	if ok {
		Self.group.Delete(host)
		atomic.AddUint32(&Self.total, ^uint32(0))
		// 这里先在map中去掉后再停止cc goroutine，防止gc后跳空指针
		cc.(*CC).Stop()
	}
}

func (Self *CCGroup) Len() int {
	return int(atomic.LoadUint32(&Self.total))
}
